#!/bin/bash
set -euo pipefail

# 脚本名称：test_disk_speed.sh

# 功能说明：
# 本脚本用于测试物理磁盘和多种格式的虚拟磁盘的读写速度。
# 主要流程包括：
# 1. 检查并安装缺失的依赖工具。
# 2. 确保 NBD 模块已加载。
# 3. 进行读写速度测试。
# 4. 处理虚拟磁盘（支持多种格式：QCOW2, raw, vmdk, vhd, vhdx, img）。
# 5. 查找可用的 NBD 设备并挂载虚拟磁盘。
# 6. 进行读写速度测试。
# 7. 输出测试结果。

# 注意事项：
# - 需要 root 权限运行部分命令，如 qemu-nbd。
# - 在执行测试前，建议备份重要数据，以防万一。

# 需要安装的依赖包映射（命令->包）
declare -A PKG_MAP=(
    ["qemu-img"]="qemu-utils"
    ["qemu-nbd"]="qemu-utils"
    ["parted"]="parted"
    ["fdisk"]="util-linux"
    ["hdparm"]="hdparm"
    ["bc"]="bc"
)

# 检查并安装缺失的软件包
install_missing_packages() {
    local missing_pkgs=()

    # 检查每个命令对应的包是否已安装
    for cmd in "${!PKG_MAP[@]}"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            pkg="${PKG_MAP[$cmd]}"
            # 避免重复添加包
            if [[ ! " ${missing_pkgs[@]} " =~ " $pkg " ]]; then
                missing_pkgs+=("$pkg")
            fi
        fi
    done

    # 如果有缺失的包，则安装
    if [[ ${#missing_pkgs[@]} -gt 0 ]]; then
        echo "检测到缺少的软件包: ${missing_pkgs[*]}"
        echo "正在安装缺失的软件包..."
        sudo apt-get update
        sudo apt-get install -y "${missing_pkgs[@]}"
    fi
}

# 确保 NBD 模块已加载
if ! lsmod | grep -q "^nbd"; then
    echo "加载 NBD 内核模块..."
    sudo modprobe nbd max_part=16
fi

# 初始化 NBD_DEVICE 变量
NBD_DEVICE=""

# 退出时清理 NBD 设备
cleanup() {
    if [[ -n "$NBD_DEVICE" ]]; then
        echo "卸载 NBD 设备: $NBD_DEVICE"
        sudo qemu-nbd --disconnect "$NBD_DEVICE"
    fi
}
trap cleanup EXIT

# 进行读写速度测试
test_disk_speed() {
    local disk="$1"
    local test_file="/tmp/testfile"

    echo "测试写入速度: $disk"
    # 写入测试
    sync
    dd if=/dev/zero of="$test_file" bs=1M count=1024 oflag=direct 2>&1 | grep "copied" | awk '{print "写入速度: " $8 $9}'

    echo "测试读取速度: $disk"
    # 读取测试
    sync
    dd if="$test_file" of=/dev/null bs=1M count=1024 iflag=direct 2>&1 | grep "copied" | awk '{print "读取速度: " $8 $9}'

    # 删除测试文件
    rm -f "$test_file"
}

# 检查输入参数
if [[ $# -ne 1 ]]; then
    echo "用法: $0 <磁盘路径或虚拟磁盘文件>"
    exit 1
fi

INPUT="$1"

# 脚本名称：test_disk_speed.sh

set -euo pipefail

# 需要安装的依赖包映射（命令->包）
declare -A PKG_MAP=(
    ["qemu-img"]="qemu-utils"
    ["qemu-nbd"]="qemu-utils"
    ["parted"]="parted"
    ["fdisk"]="util-linux"
    ["hdparm"]="hdparm"
    ["bc"]="bc"
)

# 检查并安装缺失的软件包
install_missing_packages() {
    local missing_pkgs=()

    # 检查每个命令对应的包是否已安装
    for cmd in "${!PKG_MAP[@]}"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            pkg="${PKG_MAP[$cmd]}"
            # 避免重复添加包
            if [[ ! " ${missing_pkgs[@]} " =~ " $pkg " ]]; then
                missing_pkgs+=("$pkg")
            fi
        fi
    done

    # 如果有缺失的包，则安装
    if [[ ${#missing_pkgs[@]} -gt 0 ]]; then
        echo "检测到缺少的软件包: ${missing_pkgs[*]}"
        echo "正在安装缺失的软件包..."
        sudo apt-get update
        sudo apt-get install -y "${missing_pkgs[@]}"
    fi
}

# 检查并安装缺失的依赖工具
install_missing_packages

if [[ -b "$INPUT" ]]; then
    # 物理磁盘
    test_disk_speed "$INPUT"
elif [[ -f "$INPUT" ]]; then
    # 处理虚拟磁盘（支持多种格式: QCOW2, raw, vmdk, vhd, vhdx, img）
    echo "检测到虚拟磁盘文件: $INPUT"

    # 使用 qemu-img 检查镜像格式
    IMAGE_FORMAT=$(qemu-img info "$INPUT" | grep "file format" | awk '{print $3}')
    if [[ -z "$IMAGE_FORMAT" ]]; then
        echo "错误: 无法检测到镜像格式"
        exit 1
    fi
    echo "虚拟磁盘格式: $IMAGE_FORMAT"

    # 查找可用的 NBD 设备
    for i in {0..15}; do
        if [[ ! -e "/sys/block/nbd$i/pid" ]]; then
            NBD_DEVICE="/dev/nbd$i"
            break
        fi
    done

    if [[ -z "$NBD_DEVICE" ]]; then
        echo "错误: 没有可用的 NBD 设备."
        exit 1
    fi

    echo "挂载虚拟磁盘到: $NBD_DEVICE"

    # 明确指定磁盘格式
    sudo qemu-nbd --connect="$NBD_DEVICE" --format="$IMAGE_FORMAT" "$INPUT"

    # 进行读写速度测试
    test_disk_speed "$NBD_DEVICE"

    # 退出时自动清理
else
    echo "错误: 无效的磁盘路径或文件."
    exit 1
fi

